package main

import (
	"bytes"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"github.com/kbinani/screenshot"
	"image/png"
	"io"
	"net"
	"os"
)

// Response 定义了服务器响应的结构
type Response struct {
	Code  int    `json:"code"`  // 状态码
	Error string `json:"error"` // 错误信息（如果有）
	Data  []byte `json:"data"`  // 原始图像数据
}

var address = "127.0.0.1:12345"

// sendWithLengthPrefix 使用长度前缀协议发送数据
func sendWithLengthPrefix(conn net.Conn, data []byte) error {
	var lengthBuf [4]byte
	binary.BigEndian.PutUint32(lengthBuf[:], uint32(len(data)))

	// 发送长度前缀
	_, err := conn.Write(lengthBuf[:])
	if err != nil {
		return err
	}

	// 发送实际数据
	_, err = conn.Write(data)
	return err
}

// readWithLengthPrefix 读取长度前缀协议的数据
func readWithLengthPrefix(conn net.Conn) ([]byte, error) {
	var lengthBuf [4]byte
	_, err := io.ReadFull(conn, lengthBuf[:])
	if err != nil {
		return nil, err
	}
	length := binary.BigEndian.Uint32(lengthBuf[:])

	data := make([]byte, length)
	_, err = io.ReadFull(conn, data)
	if err != nil {
		return nil, err
	}

	return data, nil
}

// captureScreen 捕获屏幕并返回 bytes.Buffer
func captureScreen() bytes.Buffer {
	bounds := screenshot.GetDisplayBounds(0)
	img, err := screenshot.CaptureRect(bounds)
	if err != nil {
		panic(err)
	}

	// 使用 bytes.Buffer 来存储图像的二进制数据
	var buf bytes.Buffer
	err = png.Encode(&buf, img)
	if err != nil {
		panic(err)
	}

	return buf
}

// sendMessage 发送 JSON 消息给服务器
func sendMessage(conn net.Conn, response Response) {
	// 序列化 Response 结构为 JSON
	jsonData, err := json.Marshal(response)
	if err != nil {
		fmt.Println("Error marshalling JSON:", err)
		return
	}

	// 发送 JSON 数据，带长度前缀
	err = sendWithLengthPrefix(conn, jsonData)
	if err != nil {
		fmt.Println("Error sending data:", err)
		return
	}

	// 可选：打印发送的 JSON 数据
	fmt.Println("Sent response:", string(jsonData))
}

// start 启动客户端连接
func start() {
	conn, err := net.Dial("tcp", address)
	if err != nil {
		fmt.Println("Error connecting:", err.Error())
		os.Exit(1)
	}
	defer conn.Close()
	fmt.Println("已连接到服务器")

	for {
		// 读取命令使用长度前缀
		commandBytes, err := readWithLengthPrefix(conn)
		if err != nil {
			if err == io.EOF {
				fmt.Println("Connection closed by server")
			} else {
				fmt.Println("Error reading command:", err.Error())
			}
			break
		}
		command := string(commandBytes)
		fmt.Printf("Received command: %s\n", command)

		response := Response{
			Code: 200,
		}

		switch command {
		case "1":
			data := captureScreen()
			response.Data = data.Bytes()
		default:
			response.Code = 400
			response.Error = "Unknown command"
		}

		// 发送 Response 作为 JSON，带长度前缀
		sendMessage(conn, response)
	}
}

func main() {
	start()
}
